import pickle
import numpy as np
import gradio as gr
from PIL import Image
import socket

# 加载模型
with open('best_knn_model.pkl', 'rb') as f:
    knn_model = pickle.load(f)
print("✅ KNN模型加载成功！")


def find_available_port(start_port=7860, max_attempts=20):
    """寻找可用的端口"""
    for port in range(start_port, start_port + max_attempts):
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.bind(('127.0.0.1', port))
                return port
        except OSError:
            continue
    return None


def predict_digit(image):
    if image is None:
        return "请在画板上绘制一个数字"

    try:
        # 确保图像数据有效
        if not isinstance(image, np.ndarray) or image.size == 0:
            return "图像数据无效，请重新绘制"

        # 转换为灰度图
        if len(image.shape) == 3:
            image_gray = np.mean(image, axis=2).astype(np.uint8)
        else:
            image_gray = image

        print(f"输入图像形状: {image_gray.shape}")

        # 颜色反转（Gradio画板是黑底白字，模型需要白底黑字）
        image_gray = 255 - image_gray

        # 调整大小到8x8
        pil_image = Image.fromarray(image_gray)
        pil_image = pil_image.resize((8, 8), Image.Resampling.LANCZOS)
        image_array = np.array(pil_image)

        # 归一化到0~16
        image_array = image_array.astype(float) / 255.0 * 16.0

        # 扁平化特征向量
        features = image_array.flatten().reshape(1, -1)

        # 预测
        pred = knn_model.predict(features)[0]
        proba = knn_model.predict_proba(features)[0]
        confidence = np.max(proba)

        return f"预测数字: {pred}\n置信度: {confidence:.3f}"

    except Exception as e:
        return f"识别失败: {str(e)}"


# 寻找可用端口
available_port = find_available_port()
if available_port:
    print(f"找到可用端口: {available_port}")

    iface = gr.Interface(
        fn=predict_digit,
        inputs=gr.Sketchpad(type="numpy"),
        outputs=gr.Textbox(label="识别结果"),
        title="手写数字识别",
        description="在画板上绘制数字0-9",
        live=True
    )

    print(f"🌐 应用将在 http://127.0.0.1:{available_port} 启动")
    iface.launch(server_name="127.0.0.1", server_port=available_port)
else:
    print("❌ 找不到可用端口，请手动关闭其他程序")